import networkx as nx
import matplotlib.pyplot as plt
from collections import defaultdict

def create_gas_network():
    # 初始化图（无向图）
    G = nx.DiGraph()

    # 添加节点（阀门、调压站、管道、用户）
    nodes = [
        (1, {"type": "Valve", "pressure": 0.4}),       # 阀门V1
        (2, {"type": "Regulator", "pressure": 0.35}),  # 调压站R1
        (3, {"type": "Pipe", "pressure": 0.3}),        # 管道P1
        (4, {"type": "Pipe", "pressure": 0.25}),       # 管道P2
        (5, {"type": "Consumer", "pressure": 0.2}),    # 用户C1
        (6, {"type": "Valve", "pressure": 0.4}),       # 阀门V2
        (7, {"type": "Regulator", "pressure": 0.35}),  # 调压站R2
        (8, {"type": "Pipe", "pressure": 0.3}),        # 管道P3
        (9, {"type": "Pipe", "pressure": 0.25}),       # 管道P4
        (10, {"type": "Consumer", "pressure": 0.2}),   # 用户C2
        (11, {"type": "Valve", "pressure": 0.4}),      # 阀门V3
        (12, {"type": "Regulator", "pressure": 0.35}), # 调压站R3
        (13, {"type": "Pipe", "pressure": 0.3}),       # 管道P5
        (14, {"type": "Consumer", "pressure": 0.2}),   # 用户C3
        (15, {"type": "Consumer", "pressure": 0.2})    # 用户C4
    ]
    G.add_nodes_from(nodes)

    # 添加边（管道连接关系，带权重）
    edges = [
        (1, 2, {"length": 500}),   # V1-R1
        (2, 3, {"length": 1000}),  # R1-P1
        (3, 4, {"length": 800}),   # P1-P2
        (4, 5, {"length": 600}),   # P2-C1
        (2, 6, {"length": 700}),   # R1-V2
        (6, 7, {"length": 900}),   # V2-R2
        (7, 8, {"length": 1200}),  # R2-P3
        (8, 9, {"length": 1000}),  # P3-P4
        (9, 10, {"length": 800}),  # P4-C2
        (2, 11, {"length": 1100}), # R1-V3
        (11, 12, {"length": 1300}),# V3-R3
        (12, 13, {"length": 1000}),# R3-P5
        (13, 14, {"length": 900}), # P5-C3
        (13, 15, {"length": 700}), # P5-C4
        (3, 12, {"length": 1500}), # P1-R3（冗余连接）
        (4, 13, {"length": 2000})  # P2-P5（冗余连接）
    ]
    G.add_edges_from((u, v, d) for u, v, d in edges)

    return G

def plot_gas_network(G, highlight_nodes=None):
    # 分离节点类型
    valves = [n for n in G.nodes if G.nodes[n]["type"] == "Valve"]
    regulators = [n for n in G.nodes if G.nodes[n]["type"] == "Regulator"]
    pipes = [n for n in G.nodes if G.nodes[n]["type"] == "Pipe"]
    consumers = [n for n in G.nodes if G.nodes[n]["type"] == "Consumer"]

    # 设置 Shell 布局
    nlist = [valves, regulators, pipes, consumers]
    pos = nx.shell_layout(G, nlist=nlist)

    # 绘制基础节点
    nx.draw_networkx_nodes(G, pos, nodelist=valves, node_color='red', node_size=600, label='Valve')
    nx.draw_networkx_nodes(G, pos, nodelist=regulators, node_color='blue', node_size=600, label='Regulator')
    nx.draw_networkx_nodes(G, pos, nodelist=pipes, node_color='green', node_size=300, label='Pipe')
    nx.draw_networkx_nodes(G, pos, nodelist=consumers, node_color='orange', node_size=400, label='Consumer')

    # 绘制边（带权重）
    edge_labels = {(u, v): d['length'] for u, v, d in G.edges(data=True)}
    nx.draw_networkx_edges(G, pos, width=2, alpha=0.6)
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8)

    # 高亮显示泄漏点和最佳关停节点
    if highlight_nodes:
        nx.draw_networkx_nodes(G, pos, nodelist=highlight_nodes, node_color='purple', node_size=800, label='Highlight')

    # 绘制标签
    nx.draw_networkx_labels(G, pos, font_size=10, font_weight='bold')

    # 显示图例和图形
    plt.legend(scatterpoints=1, borderpad=1, fontsize=10)
    plt.title("Gas Pipeline Network with MST-Based Valve Closure", fontsize=14)
    plt.axis('off')
    plt.show()

def find_best_valve_to_close(G, leak_nodes, source_node=1):
    """
    使用最小生成树（MST）方法找到最佳关停节点
    :param G: 燃气管网图（DiGraph）
    :param leak_nodes: 泄漏点节点列表
    :param source_node: 源节点（默认为节点1）
    :return: 最佳关停节点（阀门）
    """
    # 将有向图转换为无向图
    undirected_G = G.to_undirected()
    
    # 构建最小生成树（MST）
    mst = nx.minimum_spanning_tree(undirected_G, weight="length")

    # 检查源节点是否存在
    if source_node not in G.nodes:
        raise ValueError(f"源节点 {source_node} 不存在于图中")

    # 统计每个阀门在泄漏点路径中的出现次数
    valve_usage = defaultdict(int)

    for leak in leak_nodes:
        if leak not in mst.nodes:
            print(f"泄漏点 {leak} 不在MST中，跳过")
            continue
        path = nx.shortest_path(mst, source=source_node, target=leak)
        for node in path:
            if G.nodes[node]["type"] == "Valve":
                valve_usage[node] += 1

    if not valve_usage:
        raise ValueError("未找到连接泄漏点路径的阀门")

    # 选择被最多路径使用的阀门
    best_valve = max(valve_usage, key=valve_usage.get)
    return best_valve

if __name__ == "__main__":
    # 创建燃气管网
    gas_network = create_gas_network()

    # 设置多个泄漏点（例如节点13和节点5）
    leak_nodes = [13, 5]
    print(f"泄漏点位于节点 {leak_nodes}")

    # 找到最佳关停节点
    try:
        best_valve = find_best_valve_to_close(gas_network, leak_nodes)
        print(f"最佳关闭阀门为节点 {best_valve}")
    except Exception as e:
        print(f"错误: {e}")
        exit()

    # 可视化管网并高亮显示泄漏点和最佳关停节点
    plot_gas_network(gas_network, highlight_nodes=[best_valve] + leak_nodes)